import sys
import time
import asyncio
from MyTT import EMA
import numpy as np
import pandas as pd
from tqsdk import TqApi, TqAuth, TqKq
import sqlite3

async def create_api():
    """创建API实例"""
    try:
        api = TqApi(account=TqKq(), auth=TqAuth("cps168", "alibaba999"))
        # 确保API初始化完成
        await api.wait_update()
        return api
    except Exception as e:
        print(f"API创建失败: {e}")
        return None

async def run_strategy():
    """主策略函数"""
    api = await create_api()
    if api is None:
        return
        
    running = True
    active_symbols = set()
    products = get_monitored_products()
    symbols = [p[2] for p in products]
    
    try:
        await run_all_symbols(api, symbols, running, active_symbols)
        async with api.register_update_notify() as update_chan:
            async for _ in update_chan:
                if not running:
                    break
    except Exception as e:
        print(f"策略运行出错: {e}")
    finally:
        running = False
        while len(active_symbols) > 0:
            await asyncio.sleep(0.1)
        if api:
            try:
                await api.close()
            except RuntimeError as e:
                if "no running event loop" not in str(e):
                    raise

def get_monitored_products():
    """从数据库获取监控中的品种"""
    conn = sqlite3.connect('trade_monitor.db')
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM products WHERE is_monitored=1")
    products = cursor.fetchall()
    conn.close()
    return products

async def run_all_symbols(api, symbols, running, active_symbols):
    """处理所有品种"""
    print(f"开始处理品种列表: {symbols}")
    try:
        tasks = []
        for symbol in symbols:
            task = asyncio.create_task(process_product(api, symbol, running, active_symbols))
            tasks.append(task)
        await asyncio.gather(*tasks, return_exceptions=True)
    except Exception as e:
        print(f"策略执行出错: {e}")

async def process_product(api, symbol, running, active_symbols):
    """单个品种处理函数"""
    active_symbols.add(symbol)
    try:  
        quote = api.get_quote(symbol)
        klines = api.get_kline_serial(symbol, 1*60, 500)
        
        async with api.register_update_notify() as update_chan:
            async for _ in update_chan:
                if not running:
                    break
                    
                if api.is_changing(quote):
                    print(f"{time.strftime('%X')} {symbol} 行情更新，当前价: {quote.last_price}")
                    
                    q = (3*klines.close + klines.low + klines.open + klines.high) / 6
                    ma_values = calculate_ma(q)
                    short_ma = ma_values['ma']
                    ema_7 = ma_values['ema']
                    
                    position = api.get_position(symbol)
                    
                    if position.pos_short > 0 and short_ma[-2] > ema_7[-2]:
                        await api.insert_order(symbol, direction="BUY", offset="CLOSE", 
                                        volume=position.pos_short)
                        print(f"{time.strftime('%X')} 平空仓: {symbol} 手数: {position.pos_short}")
                        
                    if position.pos_long > 0 and short_ma[-2] < ema_7[-2]:
                        await api.insert_order(symbol, direction="SELL", offset="CLOSE", 
                                        volume=position.pos_long)
                        print(f"{time.strftime('%X')} 平多仓: {symbol} 手数: {position.pos_long}")

    except Exception as e:
        print(f"处理品种 {symbol} 时出错: {e}")                        
    finally:
        active_symbols.remove(symbol)

def calculate_ma(q):
    """自定义权重计算"""
    weights = np.arange(26, 0, -1)
    ma = []
    q_array = np.array(q)
    for i in range(len(q_array)):
        if i < 25:
            ma.append(np.nan)
            continue
        window = q_array[i-25:i+1]
        ma.append(np.dot(window, weights) / 351)
    
    ema_7 = EMA(pd.Series(ma), 7)
    return {'ma': ma, 'ema': ema_7.tolist()}

if __name__ == "__main__":
    print("****************************************|                        |****************************************")
    print("****************************************|    程序化移动止损程序  |****************************************")
    print("****************************************|                        |****************************************")
    print("")
    print("")
    asyncio.run(run_strategy())